import numpy as np

# 数中源代码错误
# 原代码在修改梯度时使用了临时变量grad，
# 这不会影响原列表。通过遍历索引并直接修改grads[i]，
# 可以正确更新列表的值。
def clip_grads(grads, max_norm):
    total_norm = 0
    for grad in grads:
        total_norm += grad ** 2  # 累加平方值
    total_norm = np.sqrt(total_norm)  # 计算总范数

    rate = max_norm / (total_norm + 1e-6)
    if rate < 1:
        for i in range(len(grads)):  # 通过索引修改原列表
            grads[i] *= rate


# 开始NLP

# 语料库预处理
def preprocess(text):
    text = text.lower()
    text = text.replace('.',' .')
    words = text.split()
    # print(words)
    word_to_id,id_to_word = {},{}
    for word in words:
        if word not in word_to_id:
            tmp_id = len(word_to_id)
            word_to_id[word] = tmp_id
            id_to_word[tmp_id] = word
    corpus = [word_to_id[w] for w in words]

    return corpus,word_to_id,id_to_word

def create_co_matrix(corpus,vocab_size,window_size=1):
    corpus_size = len(corpus)
    co_matrix = np.zeros((vocab_size,vocab_size),dtype=np.int32)

    for idx,word_id in enumerate(corpus):
        for i in range(1,window_size+1):
            left_idx = idx - i
            right_idx = idx + i 

            if left_idx >= 0:
                left_word_id = corpus[left_idx]
                co_matrix[word_id,left_word_id] += 1

            if right_idx < corpus_size:
                right_word_id = corpus[right_idx]
                co_matrix[word_id,right_word_id] += 1

    return co_matrix

def cos_similarity(x,y,eps=1e-8):
    nx = x/(np.sqrt(np.sum(x**2))+eps)
    ny = y/(np.sqrt(np.sum(y**2))+eps)
    return np.dot(nx,ny)

def most_similar(query,word_to_id,id_to_word,word_matrix,top=5):
    # 取出查询词
    if query not in word_to_id:
        print(f"{query} is not found")
        return 
    
    print(f"\n[query] {query}")
    query_id = word_to_id[query]
    query_vec = word_matrix[query_id]

    # 计算余弦相似度
    vocab_size = len(id_to_word)
    similarity = np.zeros(vocab_size)
    for i in range(vocab_size):
        similarity[i] = cos_similarity(query_vec,word_matrix[i])

    # 基于余弦相似度，按降序输出值
    count= 0
    for i in (-1*similarity).argsort():
        if id_to_word[i] == query:
            continue
        print(f"{id_to_word[i]}: {similarity[i]}")

        count += 1
        if count >= top:
            return 
        
def ppmi(C, verbose=False, eps=1e-8):
    M = np.zeros_like(C, dtype=np.float32)
    N = np.sum(C)
    S = np.sum(C,axis=0)
    total = C.shape[0]*C.shape[1]
    cnt = 0

    for i in range(C.shape[0]):
        for j in range(C.shape[1]):
            pmi = np.log2(C[i,j]*N/(S[i]*S[j])+eps)
            M[i,j] = max(0,pmi)

            if verbose:
                cnt += 1
                if cnt % (total//100+1)==0:
                    print('%.1f%% done' % (100*cnt/total))
    return M

def create_contexts_target(corpus,window_size=1):
    target = corpus[window_size:-window_size]
    contexts = []

    for idx in range(window_size,len(corpus)-window_size):
        cs = []
        for t in range(-window_size,window_size+1):
            if t==0:
                continue
            cs.append(corpus[idx+t])
        contexts.append(cs)
    return np.array(contexts),np.array(target)

def convert_one_hot(corpus,vocab_size):
    N = corpus.shape[0]

    if corpus.ndim == 1:        
        one_hot = np.zeros((N,vocab_size),dtype=np.int32)
        for i,word_id in enumerate(corpus):
            one_hot[i,word_id] = 1
    elif corpus.ndim == 2:
        C = corpus.shape[1]
        one_hot = np.zeros((N,C,vocab_size),dtype=np.int32)
        for idx_0,word_ids in enumerate(corpus):
            for idx_1,word_id in enumerate(word_ids):
                one_hot[idx_0,idx_1,word_id]=1
    return one_hot
    

if __name__ == '__main__':
    # grads = [3.0, 4.0]
    # clip_grads(grads, max_norm=2)
    # print(grads)  # 输出修改后的值：[1.2, 1.6]

    text = "You say goodbye and I say hello."
    corpus,word_to_id,id_to_word = preprocess(text)
    
    print(corpus)
    contexts,target = create_contexts_target(corpus,window_size=1)
    contexts = convert_one_hot(contexts,len(word_to_id))
    target =convert_one_hot(target,len(word_to_id))
    print(target)